package edu.washington.cs.cmdsynth;

import java.util.List;
import java.util.Map;
import java.util.Properties;

import edu.stanford.nlp.dcoref.CorefChain;
import edu.stanford.nlp.dcoref.CorefCoreAnnotations.CorefChainAnnotation;
import edu.stanford.nlp.ling.CoreLabel;
import edu.stanford.nlp.ling.CoreAnnotations.NamedEntityTagAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.PartOfSpeechAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.SentencesAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TextAnnotation;
import edu.stanford.nlp.ling.CoreAnnotations.TokensAnnotation;
import edu.stanford.nlp.pipeline.Annotation;
import edu.stanford.nlp.pipeline.StanfordCoreNLP;
import edu.stanford.nlp.semgraph.SemanticGraph;
import edu.stanford.nlp.semgraph.SemanticGraphCoreAnnotations.CollapsedCCProcessedDependenciesAnnotation;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.trees.TreeCoreAnnotations.TreeAnnotation;
import edu.stanford.nlp.util.CoreMap;
import edu.washington.cs.cmdsynth.utils.Utils;

public class NLPEngine {

	 public final Properties props = new Properties();
	 public final StanfordCoreNLP pipeline;
	 
	 private String text = null;
	 private Annotation document = null;
	 
	 public NLPEngine() {
		 props.put("annotators", "tokenize, ssplit, pos, lemma, ner, parse, dcoref"); 
		 pipeline = new StanfordCoreNLP(props);
	 }
	 
	 public void parseText(String text) {
		 this.text = text;
		 //create an empty Annotation just with the given text
		 document = new Annotation(text);
		// run all Annotators on this text
		 pipeline.annotate(document);
	 }
	
	 public List<CoreMap> getParsedSentences() {
		 Utils.checkNotNull(this.text);
		 List<CoreMap> sentences = document.get(SentencesAnnotation.class);
		 return sentences;
	 }
	 
	 public Map<Integer, CorefChain> getCorefGraph() {
		 Utils.checkNotNull(this.text);
		 Map<Integer, CorefChain> graph = 
		      document.get(CorefChainAnnotation.class);
		 return graph;
	 }
	 
	 //test methods
	 public void showRelations() {
		 List<CoreMap> sentences = this.getParsedSentences();
		 for(CoreMap sentence: sentences) {
		      System.out.println(sentence);
		      System.out.println(sentence.getClass());
		      // traversing the words in the current sentence
		      // a CoreLabel is a CoreMap with additional token-specific methods
		      for (CoreLabel token: sentence.get(TokensAnnotation.class)) {
		        // this is the text of the token
		        String word = token.get(TextAnnotation.class);
		        // this is the POS tag of the token
		        String pos = token.get(PartOfSpeechAnnotation.class);
		        // this is the NER label of the token
		        String ne = token.get(NamedEntityTagAnnotation.class);  
		        System.out.println(" token: " + token + ",  its id: " + token.index());
		        System.out.println(" word: " + word + ", pos: " + pos + ", NER: " + ne);
		      }

		      // this is the parse tree of the current sentence
		      Tree tree = sentence.get(TreeAnnotation.class);

		      System.out.println("Parse tree: ");
		      System.out.println(tree);
		      
		      // this is the Stanford dependency graph of the current sentence
		      SemanticGraph dependencies = sentence.get(CollapsedCCProcessedDependenciesAnnotation.class);
		      System.out.println("dependence graph: ");
		      System.out.println(dependencies);
		    }
	 }
	 
	 public void showCorefGraph() {
		 Map<Integer, CorefChain> graph = this.getCorefGraph();
		 System.out.println("Co-ref graph, size: " + graph.size());
		 for(Integer key : graph.keySet()) {
		  	System.out.println(" " + key + ", " + graph.get(key).toString());
		 }
	 }
}